from __future__ import print_function, division
import os
from PIL import Image
import numpy as np
from torch.utils.data import Dataset
from torchvision import transforms
from . import custom_transforms as tr


class lowDataSet(Dataset):
    """
    低对比度道路 训练集&测试集（暂无验证集）
    """
    NUM_CLASSES = 1

    def __init__(self,
                 args,
                 base_dir='/home/jovyan/work/Low_Contrast_Roads/',
                 split='train',
                 ):
        """
        :param base_dir: path to dataset directory
        :param split: train/test
        """
        super().__init__()
        self._base_dir = base_dir
        self._img_dir = os.path.join(self._base_dir, 'crops', 'img')  # 图像目录
        self._gt_dir = os.path.join(self._base_dir, 'crops', 'gt')  # 掩膜目录

        self.split = split

        self.args = args
        
        _splits_dir = os.path.join(self._base_dir)

        self.img_ids = []
        self.imgs = []
        self.gts = []

        # 读取训练集/测试集文件
        with open(os.path.join(self._base_dir, 'crops', self.split + '_crops.txt'), "r") as f:
            lines = f.read().splitlines()

        # 遍历训练集/测试集
        for line in lines:
            _img = os.path.join(self._img_dir, line + '_sat.jpg')
            _gt = os.path.join(self._gt_dir, line + '_mask.png')
            assert os.path.isfile(_img)
            assert os.path.isfile(_gt)
        
            self.img_ids.append(line)
            self.imgs.append(_img)
            self.gts.append(_gt)

        assert (len(self.imgs) == len(self.gts)), "{}集合，图像与掩膜的数量不相等".format(self.split)

        # Display stats
        print('Number of images in {}: {:d}'.format(self.split, len(self.imgs)))

    def __len__(self):
        return len(self.imgs) // self.args.batch_size * self.args.batch_size

    def __getitem__(self, index):
        _img = Image.open(self.imgs[index]).convert('RGB')
        _gt = Image.open(self.gts[index])
        sample = {'image': _img, 'label': _gt}

        if self.split == "train":
            return self.transform_tr(sample)
        elif self.split == 'test':
            return self.transform_test(sample)

    # 对训练集图像进行处理
    def transform_tr(self, sample):
        # new
        composed_transforms = transforms.Compose([
            tr.RandomRotate(180),
            tr.RandomHorizontalFlip(),
            tr.RandomGaussianBlur(),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])

        return composed_transforms(sample)

    # 对测试集图像进行处理
    def transform_test(self, sample):

        composed_transforms = transforms.Compose([
            tr.FixedResize(size=self.args.crop_size),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])

        return composed_transforms(sample)


if __name__ == '__main__':
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt
    import argparse

    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    args.base_size = 512
    args.crop_size = 512
    args.batch_size = 1
    args.dataset = 'DeepGlobe'

    train_ds = deepDataSet(args, split='train')
    valid_ds = deepDataSet(args, split='valid')

    train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=0)
    valid_loader = DataLoader(valid_ds, batch_size=args.batch_size, shuffle=False, num_workers=0)

    for ii, sample in enumerate(valid_loader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()
            gt = sample['label'].numpy()
            tmp = np.array(gt[jj]).astype(np.uint8)
            # segmap = decode_segmap(tmp, dataset='DeepGlobe')
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0])
            img_tmp *= (0.229, 0.224, 0.225)
            img_tmp += (0.485, 0.456, 0.406)
            img_tmp *= 255.0
            img_tmp = img_tmp.astype(np.uint8)
            plt.figure()
            plt.title('display')
            plt.subplot(211)
            plt.imshow(img_tmp)
            plt.subplot(212)
            plt.imshow(tmp)

        if ii == 1:
            break

    plt.show(block=True)


